#include <stdio.h>
#include <errno.h>

#include <sys/time.h>
#include <sys/socket.h>

#include <arpa/inet.h>

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_wifi.h"
#include "esp_event_loop.h"
#include "nvs_flash.h"
#include "esp_log.h"

#define EXAMPLE_WIFI_SSID "smroute"
#define EXAMPLE_WIFI_PASS "LXY960712honor"

#define SERVER_STR_ADDR "192.168.0.114"
#define SERVER_PORT 8366

#define SEND_BLOCK_SIZE 1024
#define SEND_BLOCK_NUM 10

#define CLIENT_HDLE_INVL_USEC 40000

#define SOCKET_WIN_SIZE 4096

#define PRINTF(fmt, ...) do{\
        struct timeval now;\
        gettimeofday(&now, NULL);\
        printf("I {line %d}[%ld.%ld]: " fmt "\n", __LINE__, now.tv_sec, now.tv_usec / 1000, ##__VA_ARGS__);\
    } while(0)
#define PERROR() printf("E {line %d}[%d]\n", __LINE__, errno);perror("      ");

static const char *TAG = "tcp_flow_ctrl";

static EventGroupHandle_t wifi_event_group;

const int IPV4_GOTIP_BIT = BIT0;
const int IPV6_GOTIP_BIT = BIT1;

static uint8_t s_recv_buff[SEND_BLOCK_SIZE] = {0};

static esp_err_t event_handler(void *ctx, system_event_t *event)
{
    switch (event->event_id) {
    case SYSTEM_EVENT_STA_START:
        esp_wifi_connect();
        ESP_LOGI(TAG, "SYSTEM_EVENT_STA_START");
        break;
    case SYSTEM_EVENT_STA_CONNECTED:
        /* enable ipv6 */
        tcpip_adapter_create_ip6_linklocal(TCPIP_ADAPTER_IF_STA);
        break;
    case SYSTEM_EVENT_STA_GOT_IP:
        xEventGroupSetBits(wifi_event_group, IPV4_GOTIP_BIT);
        ESP_LOGI(TAG, "SYSTEM_EVENT_STA_GOT_IP");
        break;
    case SYSTEM_EVENT_STA_DISCONNECTED:
        /* This is a workaround as ESP32 WiFi libs don't currently auto-reassociate. */
        esp_wifi_connect();
        xEventGroupClearBits(wifi_event_group, IPV4_GOTIP_BIT);
        xEventGroupClearBits(wifi_event_group, IPV6_GOTIP_BIT);
        break;
    case SYSTEM_EVENT_AP_STA_GOT_IP6:
        xEventGroupSetBits(wifi_event_group, IPV6_GOTIP_BIT);
        ESP_LOGI(TAG, "SYSTEM_EVENT_STA_GOT_IP6");

        char *ip6 = ip6addr_ntoa(&event->event_info.got_ip6.ip6_info.ip);
        ESP_LOGI(TAG, "IPv6: %s", ip6);
    default:
        break;
    }
    return ESP_OK;
}

static void initialise_wifi(void)
{
    tcpip_adapter_init();
    wifi_event_group = xEventGroupCreate();
    ESP_ERROR_CHECK( esp_event_loop_init(event_handler, NULL) );
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK( esp_wifi_init(&cfg) );
    ESP_ERROR_CHECK( esp_wifi_set_storage(WIFI_STORAGE_RAM) );
    wifi_config_t wifi_config = {
        .sta = {
            .ssid = EXAMPLE_WIFI_SSID,
            .password = EXAMPLE_WIFI_PASS,
        },
    };
    ESP_LOGI(TAG, "Setting WiFi configuration SSID %s...", wifi_config.sta.ssid);
    ESP_ERROR_CHECK( esp_wifi_set_mode(WIFI_MODE_STA) );
    ESP_ERROR_CHECK( esp_wifi_set_config(ESP_IF_WIFI_STA, &wifi_config) );
    ESP_ERROR_CHECK( esp_wifi_start() );
}

static void deinitialise_wifi(void)
{
    ESP_ERROR_CHECK( esp_wifi_disconnect() );
    ESP_ERROR_CHECK( esp_wifi_stop() );
    ESP_ERROR_CHECK( esp_wifi_deinit() );
}

static void wait_for_ip()
{
    uint32_t bits = IPV4_GOTIP_BIT | IPV6_GOTIP_BIT ;

    ESP_LOGI(TAG, "Waiting for AP connection...");
    xEventGroupWaitBits(wifi_event_group, bits, false, true, portMAX_DELAY);
    ESP_LOGI(TAG, "Connected to AP");
}

void app_main()
{
    ESP_ERROR_CHECK( nvs_flash_init() );
    initialise_wifi();
    wait_for_ip();

    int recv_bytes_cnt = 0;
    int send_bytes_cnt = 0;

    int client_fd = socket(AF_INET, SOCK_STREAM, 0);
    int win_size = SOCKET_WIN_SIZE;
    int ret = setsockopt(client_fd, SOL_SOCKET, SO_RCVBUF, &win_size, sizeof(win_size));
    if(0 > ret)
    {
        PERROR();
        goto fail;
    }
    // ret = setsockopt(client_fd, SOL_SOCKET, SO_SNDBUF, &win_size, sizeof(win_size));
    // if(0 > ret)
    // {
    //     PERROR();
    //     goto fail;
    // }

    struct sockaddr_in server_addr = 
    {
        .sin_family = AF_INET,
        .sin_addr.s_addr = inet_addr(SERVER_STR_ADDR),
        .sin_port = htons(SERVER_PORT),
    };
    ret = connect(client_fd, (struct sockaddr*)&server_addr, sizeof(server_addr));
    if(0 > ret)
    {
        PERROR();
        goto fail;
    }
    PRINTF("tcp flow control test client start");
    int bytes = 0;
    struct timeval start, end, intvl;
    gettimeofday(&start, NULL);
    while(1)
    {
        int bytes_in_block = 0;
        while(bytes_in_block < SEND_BLOCK_SIZE)
        {
            int bytes = recv(client_fd, s_recv_buff + bytes_in_block, SEND_BLOCK_SIZE - bytes_in_block, 0);
            if(bytes > 0)
            {
                // PRINTF("recv %d bytes once", bytes);
                bytes_in_block += bytes;
                recv_bytes_cnt += bytes;
            }
            else if(0 == bytes)
            {
                PRINTF("server close socket");
                goto success;
            }
            else
            {
                PERROR();
                goto fail;
            }
        }
        // PRINTF("recv one block");
        usleep(CLIENT_HDLE_INVL_USEC);
        bytes_in_block = 0;
        while(bytes_in_block < SEND_BLOCK_SIZE)
        {
            int bytes = send(client_fd, s_recv_buff + bytes_in_block, SEND_BLOCK_SIZE - bytes_in_block, 0);
            if(bytes > 0)
            {
                bytes_in_block += bytes;
                send_bytes_cnt += bytes;
            }
            else if(0 == bytes)
            {
                PRINTF("send 0 bytes");
            }
            else
            {
                PERROR();
                goto fail;
            }
        }
        // PRINTF("send one block");
    }

success:
    gettimeofday(&end, NULL);
    timersub(&end, &start, &intvl);
    PRINTF("sum of recv %d bytes", recv_bytes_cnt);
    PRINTF("sum of send %d bytes", send_bytes_cnt);
    PRINTF("tcp flow control test client end spend %ld.%ld Sec", intvl.tv_sec, intvl.tv_usec / 1000);
    close(client_fd);

    sleep(1);
    deinitialise_wifi();
    PRINTF("FIN");
    return;
fail:
    if(client_fd > 0)
    {
        close(client_fd);
    }

    sleep(1);
    deinitialise_wifi();
    PRINTF("FIN");
    return;
}